#pragma once
#include "tools.h"

bool readIntFromSD(File file, int &intVar)
{
	if (file.available())
	{
		String line = "";
		char ch;
		while (file.available())
		{
			ch = file.read();
			if (ch == ' ' || ch == '\n')
				break;
			line += ch;
		}
		intVar = line.toInt();
		return true;
	}
	return false;
}

bool readDoubleFromSD(File file, double &doubleVar)
{
	if (file.available())
	{
		String line = "";
		char ch;
		while (file.available())
		{
			ch = file.read();
			if (ch == ' ' || ch == '\n')
				break;
			line += ch;
		}
		doubleVar = line.toDouble();
		return true;
	}
	return false;
}

bool readWordFromSD(File file, int maxLength, char *cmd)
{
	if (file.available())
	{
		int index = 0;
		while (file.available() && index < maxLength - 1)
		{
			char ch = file.read();
			if (index == 0 && (ch == ' ' || ch == '\r' || ch == '\n'))
			{
				continue;
			}
			if (ch == ' ' || ch == '\n' || ch == '\r')
			{
				// 如果到达行尾，停止读取
				break;
			}
			cmd[index++] = ch;
		}
		cmd[index] = '\0'; // 确保字符串以空字符结尾
		return true;
	}
	return false;
}

bool readLineFromSD(File file, int maxLength, char *cmd)
{
	if (file.available())
	{
		int index = 0;
		while (file.available() && index < maxLength - 1)
		{
			char ch = file.read();
			if (ch == '\n')
			{
				// 如果到达行尾，停止读取
				cmd[index++] = '\n';
				break;
			}
			cmd[index++] = ch;
		}
		cmd[index] = '\0'; // 确保字符串以空字符结尾
		return true;
	}
	return false;
}

// 将计算特征写入文件
// 将计算特征写入文件
void writeToTxt(const char *filepath, int label, const float *character, int character_size)
{
	LOG("\n\nwriteToTxt\n");

	// 打印character数组中的内容
	for (int i = 0; i < character_size; i++)
	{
		// LOG(i, ":", character[i], " ");
		Serial.print(i);
		Serial.print(":");
		Serial.print(character[i], 7);
		Serial.print(" ");
	}
	LOG("\n");

	// 检查文件是否存在
	if (SD.exists(filepath))
	{
		LOG("文件已存在：", filepath, "\n");
		SD.remove(filepath); // 文件存在则删除
		LOG("文件已删除：", filepath, "\n");
	}

	// 以写入模式打开文件
	File dataFile = SD.open(filepath, FILE_WRITE);
	if (!dataFile)
	{
		LOG("文件打开失败。\n");
		return;
	}
	LOG("文件打开成功。\n");

	// 将文件指针移到文件开头（这一步可以省略，写入模式本身会覆盖内容）
	dataFile.seek(0);

	// 写入label和character特征值
	dataFile.print(label);
	dataFile.print(" ");
	for (int i = 0; i < character_size; i++)
	{
		dataFile.print(i);
		dataFile.print(":");
		dataFile.print(character[i], 7);
		if (i != (character_size - 1))
		{
			dataFile.print(" "); // 每个特征之间加空格
		}
	}
	dataFile.println(); // 换行

	// 关闭文件
	dataFile.close();
	LOG("数据写入文件成功！\n");
}

// 实现简单的汉宁窗函数
float *hanningWindow(int size)
{
	LOG("\n\nhanningWindow\n");

	float *window = Malloc(float, size);
	if (!window)
	{
		LOG("Failed to allocate memory for hanningWindow\n");
		return nullptr;
	}

	for (int i = 0; i < size; ++i)
	{
		window[i] = 0.5f * (1.f - std::cos(2.f * M_PI * i / (size - 1)));
	}

	LOG("hanningWindow end\n");
	return window;
}

// 实现简化的STFT函数
float ***simplifiedSTFT(const float *x, int x_size, int n_fft, int hop_length, const std::string &window, bool center, const std::string &mode)
{
	LOG("\n\nsimplifiedSTFT\n");

	float *win = nullptr;
	if (window == "hann")
	{
		win = hanningWindow(n_fft); // size: 400 * float = 1600 B = 1.6 KB
		if (!win)
		{
			LOG("Failed to create hanning window\n");
			return nullptr;
		}
	}

	// if (SD.exists("/win_stft.txt")) {
	// 	SD.remove("/win_stft.txt");
	// }
	// File win_data = SD.open("/win_stft.txt", FILE_WRITE);
	// for (int i = 0; i < n_fft; i++) {
	// 	win_data.print(win[i], 9);
	// 	win_data.print(" ");
	// }
	// win_data.println();
	// win_data.print("count:");
	// win_data.println(n_fft);
	LOG("called hanningWindow done.\n");

	// 如果需要，进行填充
	int pad_len = center ? n_fft / 2 : 0; // 200
	int padded_size = x_size + 2 * pad_len; // 1400
	// win_data.printf("pad_len:");
	// win_data.println(pad_len);
	// win_data.printf("padded_size:");
	// win_data.println(padded_size);

	float *x_padded = Malloc(float, padded_size); // 1400 * float = 5600 B = 5.5 KB
	if (!x_padded)
	{
		LOG("Failed to allocate memory for x_padded\n");
		return nullptr;
	}

	for (int i = 0; i < padded_size; ++i)
	{
		x_padded[i] = (i >= pad_len && i < pad_len + x_size) ? x[i - pad_len] : 0;
	}
	LOG("x_padded calculate done.\n");

	// for (int i = 0; i < padded_size; i++) {
	// 	win_data.print(x_padded[i], 7);
	// 	win_data.print(" ");
	// }
	// win_data.println();
	// win_data.print("count:");
	// win_data.println(padded_size);

	// 计算帧数
	int n_frames = 1 + (padded_size - n_fft) / hop_length; // 7
	LOG("n_frames: ", n_frames, "\n");
	// win_data.printf("n_frames:");
	// win_data.println(n_frames);
	// win_data.close();

	// 使用 Malloc 分配内存
	float ***X = Malloc(float **, n_frames);
	if (!X)
	{
		LOG("Failed to allocate memory for X\n");
		return nullptr;
	}

	for (int i = 0; i < n_frames; i++)
	{
		X[i] = Malloc(float *, n_fft / 2 + 1);
		if (!X[i])
		{
			LOG("Failed to allocate memory for X[i]\n");
			return nullptr;
		}

		for (int j = 0; j < n_fft / 2 + 1; j++)
		{
			X[i][j] = Malloc(float, 2); // 2 表示实部和虚部
			if (!X[i][j])
			{
				LOG("Failed to allocate memory for X[i][j]\n");
				return nullptr;
			}
		}
	}
	LOG("Malloc X done.\n");
	
	if (SD.exists("/X_stft.txt")) {
		SD.remove("/X_stft.txt");
	}
	// File X_data = SD.open("/X_stft.txt", FILE_WRITE);
	// 对每一帧进行傅立叶变换
	// x_frame完全无差异，X[i][k][0]和X[i][k][1]有差异，约为小数点后7位
	for (int i = 0; i < n_frames; ++i)
	{
		// X_data.printf("frame:");
		// X_data.println(i);

		float *x_frame = Malloc(float, n_fft);
		if (!x_frame)
		{
			LOG("Failed to allocate memory for x_frame\n");
			return nullptr;
		}
		for (int j = 0; j < n_fft; ++j)
		{
			x_frame[j] = win[j] * x_padded[i * hop_length + j];
			// X_data.print(x_frame[j], 7);
			// X_data.print(" ");
		}
		// X_data.println();
		// X_data.print("count:");
		// X_data.println(n_fft);

		// 进行傅立叶变换
		for (int k = 0; k < n_fft / 2 + 1; ++k)
		{
			float real_sum = 0;
			float imag_sum = 0;
			for (int j = 0; j < n_fft; ++j)
			{
				real_sum += x_frame[j] * std::cos(2 * M_PI * j * k / n_fft);
				imag_sum -= x_frame[j] * std::sin(2 * M_PI * j * k / n_fft);
			}
			X[i][k][0] = real_sum;
			X[i][k][1] = imag_sum;
			// X_data.print(X[i][k][0], 7);
			// X_data.print("-");
			// X_data.print(X[i][k][1], 7);
			// X_data.print(" ");
		}
		// X_data.println();
		// X_data.print("k-count:");
		// X_data.println(n_fft / 2 + 1);
		// X_data.println();
		// X_data.println();
		Free(x_frame);
	}
	// X_data.close();
	Free(x_padded);
	Free(win);

	printMemoryInfo("simplifiedSTFT end");
	return X;
}

float **melFilterBank(int n_fft, int n_mel, int sample_rate)
{
	// 计算频率对应的Mel频率
	float *mel_frequencies = Malloc(float, n_mel + 2);

	float mel_min = 0;
	float mel_max = 2595 * log10(1 + sample_rate / 2.f / 700);

	for (int i = 0; i < n_mel + 2; ++i)
	{
		mel_frequencies[i] = mel_min + i * (mel_max - mel_min) / (n_mel + 1);
	}

	// 将Mel频率转换为线性频率
	float *frequencies = Malloc(float, n_mel + 2);

	for (int i = 0; i < n_mel + 2; ++i)
	{
		frequencies[i] = 700 * (pow(10, mel_frequencies[i] / 2595) - 1);
	}

	// 初始化梅尔滤波器组
	float **filters = Malloc(float *, n_mel);
	for (int i = 0; i < n_mel; ++i)
	{
		filters[i] = Malloc(float, n_fft / 2 + 1);
		for (int j = 0; j < n_fft / 2 + 1; ++j)
		{
			filters[i][j] = 0;
		}
	}

	for (int i = 0; i < n_mel; ++i)
	{
		float mel_center = mel_frequencies[i + 1];
		float mel_left = mel_frequencies[i];
		float mel_right = mel_frequencies[i + 2];
		float f_center = frequencies[i + 1];
		float f_left = frequencies[i];
		float f_right = frequencies[i + 2];

		for (int j = 0; j <= n_fft / 2; ++j)
		{
			float freq = sample_rate * j / static_cast<float>(n_fft);
			if (freq >= f_left && freq <= f_right)
			{
				if (freq <= f_center)
					filters[i][j] = (freq - f_left) / (f_center - f_left);
				else
					filters[i][j] = (f_right - freq) / (f_right - f_center);
			}
		}
	}

	Free(mel_frequencies);
	Free(frequencies);

	return filters;
}

float **melSpectrogram(const float *x, int x_size, int sample_rate, int n_fft, int n_hop, const std::string &win, bool center, const std::string &mode, float ref_power, int n_mel, int fmin, int fmax)
{
	// 计算STFT
	float *window = hanningWindow(n_fft);
	int pad_len = center ? n_fft / 2 : 0;

	float *x_padded = Malloc(float, x_size + 2 * pad_len);
	for (int i = 0; i < x_size + 2 * pad_len; i++)
	{
		x_padded[i] = 0;
	}
	for (int i = pad_len; i < pad_len + x_size; i++)
	{
		x_padded[i] = x[i - pad_len];
	}

	int n_frames = 1 + (x_size + 2 * pad_len - n_fft) / n_hop;

	float ***stft_result = Malloc(float **, n_frames);
	for (int i = 0; i < n_frames; i++)
	{
		stft_result[i] = Malloc(float *, n_fft / 2 + 1);
		for (int j = 0; j < n_fft / 2 + 1; j++)
		{
			stft_result[i][j] = Malloc(float, 2);
		}
	}

	for (int i = 0; i < n_frames; ++i)
	{
		float *x_frame = Malloc(float, n_fft);
		for (int j = 0; j < n_fft; ++j)
		{
			x_frame[j] = window[j] * x_padded[i * n_hop + j];
		}
		for (int k = 0; k < n_fft / 2 + 1; ++k)
		{
			float real_sum = 0;
			float imag_sum = 0;
			for (int j = 0; j < n_fft; ++j)
			{
				real_sum += x_frame[j] * std::cos(2 * M_PI * j * k / n_fft);
				imag_sum -= x_frame[j] * std::sin(2 * M_PI * j * k / n_fft);
			}
			stft_result[i][k][0] = real_sum;
			stft_result[i][k][1] = imag_sum;
		}
		Free(x_frame);
	}

	// 计算梅尔滤波器组
	float **filters = melFilterBank(n_fft, n_mel, sample_rate);

	// 只选择感兴趣的频率范围内的数据
	int n_fmin_bin = static_cast<int>(std::floor(fmin * n_fft / sample_rate));
	int n_fmax_bin = static_cast<int>(std::ceil(fmax * n_fft / sample_rate)) + 1;

	float ***stft_result_selected = Malloc(float **, n_frames);
	for (int i = 0; i < n_frames; i++)
	{
		stft_result_selected[i] = Malloc(float *, n_fmax_bin - n_fmin_bin);
		for (int j = 0; j < n_fmax_bin - n_fmin_bin; j++)
		{
			stft_result_selected[i][j] = Malloc(float, 2);
		}
	}

	for (int i = 0; i < n_frames; ++i)
	{
		for (int k = 0; k < n_fmax_bin - n_fmin_bin; ++k)
		{
			stft_result_selected[i][k][0] = stft_result[i][n_fmin_bin + k][0];
			stft_result_selected[i][k][1] = stft_result[i][n_fmin_bin + k][1];
		}
	}

	// 计算梅尔频谱图
	float **mel_spectrogram = Malloc(float *, n_frames);
	for (int i = 0; i < n_frames; i++)
	{
		mel_spectrogram[i] = Malloc(float, n_mel);
		for (int j = 0; j < n_mel; j++)
		{
			mel_spectrogram[i][j] = 0;
		}
	}

	for (int i = 0; i < n_frames; ++i)
	{
		for (int j = 0; j < n_mel; ++j)
		{
			float sum = 0;
			for (int k = 0; k < n_fft / 2 + 1; ++k)
			{
				float magnitude = std::sqrt(stft_result[i][k][0] * stft_result[i][k][0] + stft_result[i][k][1] * stft_result[i][k][1]);
				sum += magnitude * filters[j][k];
			}
			mel_spectrogram[i][j] = sum;
		}
	}

	// free memory
	Free(x_padded);
	for (int i = 0; i < n_frames; ++i)
	{
		for (int j = 0; j < n_fft / 2 + 1; ++j)
		{
			Free(stft_result[i][j]);
		}
		Free(stft_result[i]);
	}
	Free(stft_result);

	for (int i = 0; i < n_frames; ++i)
	{
		for (int j = 0; j < n_fmax_bin - n_fmin_bin; ++j)
		{
			Free(stft_result_selected[i][j]);
		}
		Free(stft_result_selected[i]);
	}
	Free(stft_result_selected);

	for (int i = 0; i < n_mel; ++i)
	{
		Free(filters[i]);
	}
	Free(filters);

	return mel_spectrogram;
}